import torch
import torch.nn as nn

class LBNet(nn.Module):

    def __init__(self):

        super(LBNet, self).__init__()
        self.convolutions = nn.Sequential(
            nn.Conv2d(2, 16, kernel_size=7, stride=1),
            nn.ReLU(),
            nn.LocalResponseNorm(5, 0.0001, 0.75, 2),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(16, 64, kernel_size=7, stride=1),
            nn.ReLU(),
            nn.LocalResponseNorm(5, 0.0001, 0.75, 2),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(64, 256, kernel_size=7, stride=1)
        )
        self.mlp = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(21 * 11 * 256,1),
            nn.Sigmoid()
        )

    def forward(self, x,flag):
        if flag == 'train':
            x1 = x[0,:,:,:]
            x2 = x[1, :, :, :]
            x3 = x[2, :, :, :]
            x4 = x[3, :, :, :]
            x5 = x[4, :, :, :]
            x6 = x[5, :, :, :]
            x7 = x[6, :, :, :]
            x8 = x[7, :, :, :]
            x9 = x[8, :, :, :]
            x10 = x[9, :, :, :]
            x11 = x[10, :, :, :]
            x12 = x[11, :, :, :]
            x13 = x[12, :, :, :]
            x14 = x[13, :, :, :]
            x15 = x[14, :, :, :]
            x16 = x[15, :, :, :]
            x17 = x[16, :, :, :]
            x18 = x[17, :, :, :]
            x19 = x[18, :, :, :]
            x20 = x[19, :, :, :]
            positive_pairs = [torch.cat((x1,x2),0),torch.cat((x3,x4),0),torch.cat((x5,x6),0),torch.cat((x7,x8),0),torch.cat((x9,x10),0),
                              torch.cat((x11,x12),0),torch.cat((x13,x14),0),torch.cat((x15,x16),0),torch.cat((x17,x18),0),torch.cat((x19,x20),0)]
            negative_pairs = [torch.cat((x1,x11),0),torch.cat((x2,x12),0),torch.cat((x3,x13),0),torch.cat((x4,x14),0),torch.cat((x5,x15),0),
                              torch.cat((x6,x16),0),torch.cat((x7,x17),0),torch.cat((x8,x18),0),torch.cat((x9,x19),0),torch.cat((x10,x20),0)]
            pairs = positive_pairs + negative_pairs
            pairs = [i.unsqueeze(0) for i in pairs]
            pairs = torch.cat(pairs,0)

            x = self.convolutions(pairs)
            x = x.view(-1, 21 * 11 * 256)
            x = self.mlp(x)
            return x
        else:
            x = self.convolutions(x)
            x = x.view(-1, 21 * 11 * 256)
            x = self.mlp(x)
            return x

